import torch
from train import *
import TitleChatGPT
import baidu

def find_highest_scoring_title(titles, score_title):
    highest_score = float('-inf')
    highest_scoring_title = None
    print("开始训练")
    model, pca, scaler = train(8)
    print("结束训练")
    for title in titles:
        score = score_title(title, model, pca, scaler)
        if score > highest_score:
            highest_score = score
            highest_scoring_title = title

    return highest_scoring_title, highest_score


# 预测新标题的阅读量吸引度
def predict_title(title, model, pca, scaler):
    # 生成特征向量
    test_X = pre_set_title(title)

    # 将特征向量转换为numpy数组，并应用PCA和标准化
    test_X_np = test_X.numpy().reshape(1, -1)
    test_X_pca = pca.transform(test_X_np)
    test_X_scaled = scaler.transform(test_X_pca)

    # 转换回Tensor
    test_X_tensor = torch.tensor(test_X_scaled, dtype=torch.float32)

    # 预测
    predicted_read_count = model(test_X_tensor)
    return predicted_read_count.item()


# # Press the green button in the gutter to run the script.
# if __name__ == '__main__':

#     model = train(8)
#     predict = predict_title("中山大学新晋网红打卡地！",model)
#     # predict = predict_title("首聚深圳校区！中大人研讨如何推进科技自立自强",model)
#     # 决赛圈！20人！恭喜！
#     # 中山大学新晋网红打卡地！??????
#     print(predict)
def generate_title_by_model(model_name):
    results = {}
    if model_name == 'ChatGPT':
        results = TitleChatGPT.Generate_Result_by_ChatGPT()
    if model_name == 'baidu':
        results = baidu.ask_baidu()
    return results


if __name__ == '__main__':
    # model, pca, scaler = train(8)
    # predict = predict_title("中山大学新晋网红打卡地！??????", model, pca, scaler)
    # print(predict)
    # 文本输入通过content_text.txt

    result = generate_title_by_model('ChatGPT')
    # top_titles 是之前生成的10个标题列表
    top_titles = result['top_titles']

    # 找出分数最高的标题
    highest_scoring_title, highest_score = find_highest_scoring_title(top_titles, predict_title)
    print(f"Highest Scoring Title: {highest_scoring_title}")
    print(f"Score: {highest_score}")
